# 本脚本未完成编辑
import math

import numpy as np
import pandas as pd


def main(cy_sj, pd_sj, tj, qs_year, js_year):  # 主要逻辑函数，根据规则交易，并计算收益率
    # 参数含义：持有时长、判断时长、条件、起始、结束
    # 参数tj为二元列表，[买入时yj的下限,买入时yj的上限]
    # 参数qs_year表示研究起始年份,js_year表示研究结束年份，都包含端点
    df = pd.read_excel('../data/创业板指研究.xlsx', sheet_name='Sheet1')
    df['yj0'] = (df['涨幅(%)'] + 0.8 * df['涨幅(%)'].shift(1) + 0.6 * df['涨幅(%)'].shift(2)
                 + 0.4 * df['涨幅(%)'].shift(3) + 0.2 * df['涨幅(%)'].shift(4))
    df.loc[0:4, 'yj0'] = 0
    df['yj1'] = (df['涨幅(%)'].shift(1) + 0.8 * df['涨幅(%)'].shift(2) + 0.6 * df['涨幅(%)'].shift(3)
                 + 0.4 * df['涨幅(%)'].shift(4) + 0.2 * df['涨幅(%)'].shift(5))
    df.loc[0:5, 'yj1'] = 0
    df['worst'] = (df['收盘价'].shift(-cy_sj) - df['收盘价']) / df['收盘价'] * 100
    df.loc[df['worst'].isna().tail(cy_sj).index, 'worst'] = 0  # 清除最后的缺失值
    yj = df['yj1'].values  # 投资判断依据（表列转化为列表）
    spj = df['收盘价'].values  # 指数收盘价
    year, month, day = df[['year', 'month', 'day']].to_numpy().T  # 将日期转换为列表，后续会用到
    # 最关键的一行，将不满足条件的日子从投资列表中剔除
    btr = [i for i, value in enumerate(yj) if value < tj[0] or value > tj[1] or df['涨幅(%)'][i] > 1.5]  # 表示不投入的交易日
    n = 0  # 初始化——投资失败次数
    days = []  # 初始化——单个投资周期天数
    s = []  # 初始化——每次投资收益率
    tz_date = []  # 初始化——每次投资的日期
    yj_m = [100] * 4 + [5] * 3 + [5] * 3 + [-1] * (cy_sj - 10)  # 卖出时判断依据的涨幅（%）,可更改
    for i in range(len(spj) - cy_sj):
        if (i in btr) or (year[i] < qs_year or year[i] > js_year):  # 去除早于起始研究年限或晚于结束研究年限的，以及按规则不投资的数据
            continue
        for j in range(1, cy_sj + 1):
            if spj[i + j] > spj[i] * (1 + yj_m[j - 1] / 100) and spj[i + j] > spj[i + j - 1]:  # 是否卖出依据浮盈是否超过yj_m对应预设确定
                if j <= 4:
                    s.append((spj[i + j] * 0.985 / spj[i] - 1) * 100)  # 在前4个交易日内卖出会收1.5%的手续费
                else:
                    s.append((spj[i + j] / spj[i] - 1) * 100)  # 一周后卖出无手续费
                tz_date.append([year[i], month[i], day[i]])  # 记录盈利的投资日期
                days.append(j)
                break
            if j == cy_sj:
                s.append((spj[i + cy_sj] / spj[i] - 1) * 100)  # 两个月内未能盈利，按第41天价格卖出
                tz_date.append([year[i], month[i], day[i]])  # 记录失败的投资日期
                n += 1
                days.append(cy_sj)
    print(f'持有时间：{cy_sj}\t判断依据：{pd_sj} {tj}\t投资次数：{len(s)}')
    print(f'胜率：{(1 - n / len(s)) * 100:.2f}%')
    print(f'平均收益率（算数平均）：{np.mean(s):.2f}%')
    jh_pj = (math.exp(np.mean([math.log(s[i] / 100 + 1, math.e) for i in range(len(s))])) - 1) * 100
    print(f'平均收益率（几何平均）：{jh_pj:.2f}%')
    print(f'理想情况总收益率：顺投：{((jh_pj / 100 + 1) ** len(s) - 1) * 100:.2f}% 等额：{np.sum(s):.2f}%')
    print(f'平均周转天数：{np.mean(days):.2f}')


if __name__ == '__main__':
    qs_year = 2020  # 起始年份（含该年）
    js_year = 2025  # 结束年份（含该年）
    cy_sj, pd_sj = 75, 5
    tj = [3, 7.5]
    main(cy_sj, pd_sj, tj, qs_year, js_year)
